{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "%matplotlib inline"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "\nDeploy Single Shot Multibox Detector(SSD) model\n===============================================\n**Author**: `Yao Wang <https://github.com/kevinthesun>`_\n`Leyuan Wang <https://github.com/Laurawly>`_\n\nThis article is an introductory tutorial to deploy SSD models with TVM.\nWe will use GluonCV pre-trained SSD model and convert it to Relay IR\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "import tvm\nfrom tvm import te\n\nfrom matplotlib import pyplot as plt\nfrom tvm import relay\nfrom tvm.contrib import graph_executor\nfrom tvm.contrib.download import download_testdata\nfrom gluoncv import model_zoo, data, utils"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Preliminary and Set parameters\n------------------------------\n<div class=\"alert alert-info\"><h4>Note</h4><p>We support compiling SSD on both CPUs and GPUs now.\n\n  To get best inference performance on CPU, change\n  target argument according to your device and\n  follow the `tune_relay_x86` to tune x86 CPU and\n  `tune_relay_arm` for arm CPU.\n\n  To get best inference performance on Intel graphics,\n  change target argument to :code:`opencl -device=intel_graphics`.\n  But when using Intel graphics on Mac, target needs to\n  be set to `opencl` only for the reason that Intel subgroup\n  extension is not supported on Mac.\n\n  To get best inference performance on CUDA-based GPUs,\n  change the target argument to :code:`cuda`; and for\n  OPENCL-based GPUs, change target argument to\n  :code:`opencl` followed by device argument according\n  to your device.</p></div>\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "supported_model = [\n    \"ssd_512_resnet50_v1_voc\",\n    \"ssd_512_resnet50_v1_coco\",\n    \"ssd_512_resnet101_v2_voc\",\n    \"ssd_512_mobilenet1.0_voc\",\n    \"ssd_512_mobilenet1.0_coco\",\n    \"ssd_300_vgg16_atrous_voc\" \"ssd_512_vgg16_atrous_coco\",\n]\n\nmodel_name = supported_model[0]\ndshape = (1, 3, 512, 512)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Download and pre-process demo image\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "im_fname = download_testdata(\n    \"https://github.com/dmlc/web-data/blob/main/\" + \"gluoncv/detection/street_small.jpg?raw=true\",\n    \"street_small.jpg\",\n    module=\"data\",\n)\nx, img = data.transforms.presets.ssd.load_test(im_fname, short=512)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Convert and compile model for CPU.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "block = model_zoo.get_model(model_name, pretrained=True)\n\n\ndef build(target):\n    mod, params = relay.frontend.from_mxnet(block, {\"data\": dshape})\n    with tvm.transform.PassContext(opt_level=3):\n        lib = relay.build(mod, target, params=params)\n    return lib"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Create TVM runtime and do inference\n<div class=\"alert alert-info\"><h4>Note</h4><p>Use target = \"cuda -libs\" to enable thrust based sort, if you\n  enabled thrust during cmake by -DUSE_THRUST=ON.</p></div>\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "def run(lib, dev):\n    # Build TVM runtime\n    m = graph_executor.GraphModule(lib[\"default\"](dev))\n    tvm_input = tvm.nd.array(x.asnumpy(), device=dev)\n    m.set_input(\"data\", tvm_input)\n    # execute\n    m.run()\n    # get outputs\n    class_IDs, scores, bounding_boxs = m.get_output(0), m.get_output(1), m.get_output(2)\n    return class_IDs, scores, bounding_boxs\n\n\nfor target in [\"llvm\", \"cuda\"]:\n    dev = tvm.device(target, 0)\n    if dev.exist:\n        lib = build(target)\n        class_IDs, scores, bounding_boxs = run(lib, dev)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Display result\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "ax = utils.viz.plot_bbox(\n    img,\n    bounding_boxs.numpy()[0],\n    scores.numpy()[0],\n    class_IDs.numpy()[0],\n    class_names=block.classes,\n)\nplt.show()"
      ]
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "Python 3",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.6.9"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}